# encoding=utf-8
from __future__ import print_function
from tqdm import tqdm
import numpy as np
import pandas as pd
import json

def iou(box1, box2):
    area = ((box1[2] - box1[0] + 1) * (box1[3] - box1[1] + 1)) + ((box2[2] - box2[0] + 1) * (box2[3] - box2[1] + 1))
    # 计算重叠区域的左上与右下坐标
    xx1 = np.maximum(box1[0], box2[0])
    yy1 = np.maximum(box1[1], box2[1])
    xx2 = np.minimum(box1[2], box2[2])
    yy2 = np.minimum(box1[3], box2[3])
    # 计算重叠区域的长宽
    w = np.maximum(0, xx2 - xx1 + 1)
    h = np.maximum(0, yy2 - yy1 + 1)
    # 计算iou值
    iou = (w * h) / (area - (w * h))
    assert iou >= 0
    return iou

def json_average_nms(annos_file):
    defc_list1 = pd.read_json(open(annos_file), 'r')
    dfct_num1 = len(defc_list1)
    print(dfct_num1)
    img_reserved = []
    name_list = defc_list1["name"].unique()
    print(len(name_list))
    for img_name in tqdm(name_list):
        img_annos = defc_list1[defc_list1["name"] == img_name]


        img_temp_t = []
        for i in range(len(img_annos)):
            flag = 1
            for j in range(len(img_temp_t)):
                iou_value = iou(img_temp_t[j - 1]["bbox"], img_annos.iloc[i]["bbox"])
                # if iou_value > 0.8:
                if iou_value > 0.5 and img_temp_t[j - 1]["category"] == img_annos.iloc[i]["category"]:
                    if img_temp_t[j - 1]["score"] < img_annos.iloc[i]["score"]:
                        img_annos.iloc[i]["score"] = img_temp_t[j - 1]["score"] + img_annos.iloc[i]["score"] / 2
                        rect1 = img_annos.iloc[i]["bbox"]
                        rect2 = img_temp_t[j - 1]["bbox"]
                        x1 = (rect1[0] + rect2[0]) / 2
                        y1 = (rect1[1] + rect2[1]) / 2
                        x2 = (rect1[2] + rect2[2]) / 2
                        y2 = (rect1[3] + rect2[3]) / 2
                        rect = [x1, y1, x2, y2]
                        img_annos.iloc[i]["bbox"] = rect
                        img_temp_t[j - 1] = img_annos.iloc[i]
                    flag = 0
                    break
            if flag == 1:
                if img_annos.iloc[i]["score"] > 0.1:
                    img_temp_t.append(img_annos.iloc[i])
        for k, item in enumerate(img_temp_t):
            img_temp = {}
            img_temp["name"] = item["name"]
            img_temp["bbox"] = item["bbox"]
            img_temp["category"] = item["category"]
            img_temp["score"] = item["score"]
            img_reserved.append(img_temp)
    print("the numbers of reserved defects is :", len(img_reserved))
    return img_reserved

if __name__=='__main__':
    
    json_path = "../../submit/result20190920.json"
    all_dfct = json_average_nms(json_path)
    
    json.dump(all_dfct, open('../../submit/result20190920Ave.json', "wt"))


